package com.ternnework.apm.agent;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;

import org.apache.bcel.Constants;
import org.apache.bcel.classfile.ClassParser;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.ClassGen;
import org.apache.bcel.generic.ConstantPoolGen;
import org.apache.bcel.generic.InstructionConstants;
import org.apache.bcel.generic.InstructionFactory;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.MethodGen;
import org.apache.bcel.generic.ObjectType;
import org.apache.bcel.generic.PUSH;
import org.apache.bcel.generic.Type;

public class ApmClassFileTransformer implements ClassFileTransformer{
	
	
	
	public byte[] transform(ClassLoader loader, String className, Class cls,java.security.ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
	        try {
	            ClassParser classParser = new ClassParser(new java.io.ByteArrayInputStream(classfileBuffer), className + ".java"); 
	            JavaClass javaClass = classParser.parse();
	            ClassGen classGen = new ClassGen(javaClass);
	            Method[] methods = javaClass.getMethods();
	            
	            for(Method method:methods){
	            	if(method.getName().equals("testPrint")){
	            		//System.out.println(loader.ge);
	            		System.out.println(classParser.getClass().getClassLoader().getClass().getName());
	            		addTimer(classGen, method);
	            	}
	            	
	            	 //
	            }
	            ByteArrayOutputStream bos = new ByteArrayOutputStream();
	            classGen.getJavaClass().dump(bos);
	            return bos.toByteArray();
	           
	        } catch (IOException e) {
	            System.err.println(e);
	            System.exit(0);
	        }
	        return null; // No transformation required
	    }

	    private static void addTimer(ClassGen classGen, Method method) {

	        // set up the construction tools
	        InstructionFactory instructionFactory = new InstructionFactory(classGen);
	        InstructionList instructionList = new InstructionList();
	        ConstantPoolGen constantPoolGen = classGen.getConstantPool();
	        String cname = classGen.getClassName();
	        MethodGen wrapgen = new MethodGen(method, cname, constantPoolGen);
	        wrapgen.setInstructionList(instructionList);

	        // rename a copy of the original method
	        MethodGen methgen = new MethodGen(method, cname, constantPoolGen);
	        classGen.removeMethod(method);
	        String iname = methgen.getName() + "  ";
	        methgen.setName(iname);
	        classGen.addMethod(methgen.getMethod());
	        Type result = methgen.getReturnType();

	        // compute the size of the calling parameters
	        Type[] parameters = methgen.getArgumentTypes();
	        int stackIndex = methgen.isStatic() ? 0 : 1;
	        for (int i = 0; i < parameters.length; i++) {
	            stackIndex += parameters[i].getSize();
	        }

	        // save time prior to invocation
	        instructionList.append(instructionFactory.createInvoke("java.lang.System","currentTimeMillis", Type.LONG, Type.NO_ARGS,Constants.INVOKESTATIC));
	        instructionList.append(InstructionFactory.createStore(Type.LONG, stackIndex));

	        // call the wrapped method
	        int offset = 0;
	        short invoke = Constants.INVOKESTATIC;
	        if (!methgen.isStatic()) {
	            instructionList.append(InstructionFactory.createLoad(Type.OBJECT, 0));
	            offset = 1;
	            invoke = Constants.INVOKEVIRTUAL;
	        }
	        for (int i = 0; i < parameters.length; i++) {
	            Type type = parameters[i];
	            instructionList.append(InstructionFactory.createLoad(type, offset));
	            offset += type.getSize();
	        }
	        instructionList.append(instructionFactory.createInvoke(cname, iname, result, parameters, invoke));

	        // store result for return later
	        if (result != Type.VOID) {
	            instructionList.append(InstructionFactory.createStore(result, stackIndex+2));
	        }

	        // print time required for method call
	        instructionList.append(instructionFactory.createFieldAccess("java.lang.System","out",new ObjectType("java.io.PrintStream"),Constants.GETSTATIC));
	        instructionList.append(InstructionConstants.DUP);
	        instructionList.append(InstructionConstants.DUP);
	        String text = classGen.getClassName()+"."+methgen.getName() +" ";
	        instructionList.append(new PUSH(constantPoolGen, text));
	        instructionList.append(instructionFactory.createInvoke("java.io.PrintStream","print", Type.VOID, new Type[] { Type.STRING },Constants.INVOKEVIRTUAL));
	        instructionList.append(instructionFactory.createInvoke("java.lang.System", "currentTimeMillis", Type.LONG, Type.NO_ARGS, Constants.INVOKESTATIC));
	        instructionList.append(InstructionFactory.createLoad(Type.LONG, stackIndex));
	        instructionList.append(InstructionConstants.LSUB);
	        instructionList.append(instructionFactory.createInvoke("java.io.PrintStream","print", Type.VOID, new Type[] { Type.LONG },Constants.INVOKEVIRTUAL));
	        instructionList.append(new PUSH(constantPoolGen, " ms."));
	        instructionList.append(instructionFactory.createInvoke("java.io.PrintStream","println", Type.VOID, new Type[] { Type.STRING },Constants.INVOKEVIRTUAL));

	        // return result from wrapped method call
	        if (result != Type.VOID) {
	            instructionList.append(InstructionFactory.createLoad(result, stackIndex+2));
	        }
	        instructionList.append(InstructionFactory.createReturn(result));

	        // finalize the constructed method
	        wrapgen.stripAttributes(true);
	        wrapgen.setMaxStack();
	        wrapgen.setMaxLocals();
	        classGen.addMethod(wrapgen.getMethod());
	        instructionList.dispose();
	    }

}
